﻿/***************************************************************************
 * 
 * Copyright (c) 2022 str2num.com, Inc. All Rights Reserved
 * $Id$ 
 * 
 **************************************************************************/
 
 
 
/**
 * @file dtls_srtp_transport.cpp
 * @author str2num
 * @version $Revision$ 
 * @brief 
 *  
 **/

#include "ice/dtls/dtls_srtp_transport.h"

#include <rtc_base/logging.h>

#include "ice/dtls/dtls_transport_channel.h"
#include "ice/dtls/rtp_utils.h"

namespace ice {

// rfc5764
static char kDtlsSrtpExporterLabel[] = "EXTRACTOR-dtls_srtp";

DtlsSrtpTransport::DtlsSrtpTransport(const std::string& transport_name) :
        transport_name_(transport_name)
{
}

void DtlsSrtpTransport::SetDtlsTransports(DtlsTransportChannel* rtp_dtls_transport,
        DtlsTransportChannel* rtcp_dtls_transport)
{
    rtp_dtls_transport_ = rtp_dtls_transport;
    rtcp_dtls_transport_ = rtcp_dtls_transport;
    
    if (rtp_dtls_transport_) {
        rtp_dtls_transport_->SignalDtlsState.connect(this,
                &DtlsSrtpTransport::OnDtlsState);
        rtp_dtls_transport_->SignalReadPacket.connect(this,
                &DtlsSrtpTransport::OnReadPacket);
    }

    if (rtcp_dtls_transport_) {
        rtcp_dtls_transport_->SignalDtlsState.connect(this,
            &DtlsSrtpTransport::OnDtlsState);
        rtcp_dtls_transport_->SignalReadPacket.connect(this,
            &DtlsSrtpTransport::OnReadPacket);
    }

    MaybeSetupDtlsSrtp();
}

void DtlsSrtpTransport::OnDtlsState(DtlsTransportChannel* /*dtls*/,
        DtlsTransportState state)
{
    if (state != DtlsTransportState::kConnected) {
        ResetParams();
        return;
    }

    MaybeSetupDtlsSrtp();
}

void DtlsSrtpTransport::OnReadPacket(DtlsTransportChannel* /*dtls*/,
        const char* data, size_t len, int64_t ts)
{
    auto array_view = rtc::MakeArrayView(data, len);
    RtpPacketType packet_type = InferRtpPacketType(array_view);

    if (packet_type == RtpPacketType::kUnknown) {
        return;
    }

    rtc::CopyOnWriteBuffer packet(data, len);
    if (packet_type == RtpPacketType::kRtcp) {
        OnRtcpPacketReceived(std::move(packet), ts);
    } else {
        OnRtpPacketReceived(std::move(packet), ts);
    }
}

void DtlsSrtpTransport::OnRtpPacketReceived(rtc::CopyOnWriteBuffer packet,
        int64_t ts)
{
    if (!IsSrtpActive()) {
        RTC_LOG(LS_WARNING) << "Inactive SRTP transport received a rtp packet, drop it.";
        return;
    }

    char* data = packet.MutableData<char>();
    int len = (int)packet.size();
    if (!UnprotectRtp(data, len, &len)) {
        const int kFailLog = 100;
        if (unprotect_fail_count_ % kFailLog == 0) {
            RTC_LOG(LS_WARNING) << "Failed to unprotect rtp packet: "
                << ", size=" << len
                << ", seqnum=" << ParseRtpSequenceNumber(packet)
                << ", ssrc=" << ParseRtpSsrc(packet)
                << ", unprotect_fail_count=" << unprotect_fail_count_;
        }
        unprotect_fail_count_++;
        return;
    }

    packet.SetSize(len);
    SignalUnSrtpPacketReceived(this, &packet, ts);
}

void DtlsSrtpTransport::OnRtcpPacketReceived(rtc::CopyOnWriteBuffer packet,
        int64_t ts)
{
    if (!IsSrtpActive()) {
        RTC_LOG(LS_WARNING) << "Inactive SRTP transport received a rtcp packet, drop it.";
        return;
    }

    char* data = packet.MutableData<char>();
    int len = (int)packet.size();
    if (!UnprotectRtcp(data, len, &len)) {
        int type = 0;
        GetRtcpType(data, len, &type);
        RTC_LOG(LS_WARNING) << "Failed to unprotect rtcp packet: "
            << "size=" << len
            << ", type=" << type;
        return;
    }

    packet.SetSize(len);
    SignalUnSrtcpPacketReceived(this, &packet, ts);
}

int DtlsSrtpTransport::SendSrtp(const char* buf, size_t size) {
    if (!IsSrtpActive()) {
        RTC_LOG(LS_WARNING) << "Failed to send rtp packet: Inactive srtp transport";
        return -1;
    }
   
    int rtp_auth_tag_len = 0;
    GetSendAuthTagLen(&rtp_auth_tag_len, nullptr);
    rtc::CopyOnWriteBuffer packet(buf, size, size + rtp_auth_tag_len); 
    
    char* data = (char*)packet.MutableData();
    int len = (int)packet.size();
    uint16_t seq_num = ParseRtpSequenceNumber(packet);
    if (!ProtectRtp(data, len, (int)packet.capacity(), &len)) {
        RTC_LOG(LS_WARNING) << "Failed to protect rtp packet, size=" << len
            << ", seqnum=" << seq_num
            << ", ssrc=" << ParseRtpSsrc(packet)
            << ", last_send_seq_num=" << last_send_seq_num_;
        return -1;
    }
    
    last_send_seq_num_ = seq_num; 
    
    packet.SetSize(len);
    return rtp_dtls_transport_->SendPacket((const char*)packet.cdata(), packet.size());
}

int DtlsSrtpTransport::SendSrtcp(const char* buf, size_t size) {
    if (!IsSrtpActive()) {
        RTC_LOG(LS_WARNING) << "Failed to send rtcp packet: Inactive srtp transport";
        return -1;
    }
   
    int rtcp_auth_tag_len = 0;
    GetSendAuthTagLen(nullptr, &rtcp_auth_tag_len);
    rtc::CopyOnWriteBuffer packet(buf, size, size + rtcp_auth_tag_len + sizeof(uint32_t)); 
    
    char* data = (char*)packet.MutableData();
    int len = (int)packet.size();
    if (!ProtectRtcp(data, len, (int)packet.capacity(), &len)) {
        int type = 0;
        GetRtcpType(data, len, &type);
        RTC_LOG(LS_WARNING) << "Failed to protect rtcp packet, size=" << len
            << ", type=" << type;
        return -1;
    }
     
    packet.SetSize(len);
    if (rtcp_dtls_transport_) {
        return rtcp_dtls_transport_->SendPacket((const char*)packet.cdata(), packet.size());
    }
    return rtp_dtls_transport_->SendPacket((const char*)packet.cdata(), packet.size());
}

bool DtlsSrtpTransport::IsDtlsWritable() {
    return rtp_dtls_transport_ && rtp_dtls_transport_->writable() &&
        (!rtcp_dtls_transport_ || rtcp_dtls_transport_->writable());
}

void DtlsSrtpTransport::MaybeSetupDtlsSrtp() {
    if (IsSrtpActive() || !IsDtlsWritable()) {
        return;
    }

    SetupDtlsSrtp();

    if (rtcp_dtls_transport_) {
        SetupRtcpDtlsSrtp();
    }
}

void DtlsSrtpTransport::SetupDtlsSrtp() {
    std::vector<int> send_extension_ids;
    std::vector<int> recv_extension_ids;

    int selected_crypto_suite;
    rtc::ZeroOnFreeBuffer<unsigned char> send_key;
    rtc::ZeroOnFreeBuffer<unsigned char> recv_key;

    if (!ExtractParams(rtp_dtls_transport_, &selected_crypto_suite,
                &send_key, &recv_key) ||
            !SetRtpParams(selected_crypto_suite,
                &send_key[0], send_key.size(), send_extension_ids,
                selected_crypto_suite,
                &recv_key[0], recv_key.size(), recv_extension_ids))
    {
        RTC_LOG(LS_WARNING) << "DTLS-SRTP rtp param install failed";
    }
}

void DtlsSrtpTransport::SetupRtcpDtlsSrtp() {
    // Return if the DTLS-SRTP is active because the encrypted header extension
    // IDs don't need to be updated for RTCP and the crypto params don't need to
    // be reset.
    if (IsSrtpActive()) {
        return;
    }

    std::vector<int> send_extension_ids;
    std::vector<int> recv_extension_ids;
   
    int selected_crypto_suite;
    rtc::ZeroOnFreeBuffer<unsigned char> rtcp_send_key;
    rtc::ZeroOnFreeBuffer<unsigned char> rtcp_recv_key;
    if (!ExtractParams(rtcp_dtls_transport_, &selected_crypto_suite,
        &rtcp_send_key, &rtcp_recv_key) ||
        !SetRtcpParams(selected_crypto_suite, &rtcp_send_key[0],
            static_cast<int>(rtcp_send_key.size()), send_extension_ids,
            selected_crypto_suite, &rtcp_recv_key[0],
            static_cast<int>(rtcp_recv_key.size()),
            recv_extension_ids)) {
        RTC_LOG(LS_WARNING) << "DTLS-SRTP key installation for RTCP failed";
    }
}

bool DtlsSrtpTransport::ExtractParams(DtlsTransportChannel* dtls_transport,
        int* selected_crypto_suite,
        rtc::ZeroOnFreeBuffer<unsigned char>* send_key,
        rtc::ZeroOnFreeBuffer<unsigned char>* recv_key)
{
    if (!dtls_transport || !dtls_transport->IsDtlsActive()) {
        return false;
    }

    if (!dtls_transport->GetSrtpCryptoSuite(selected_crypto_suite)) {
        RTC_LOG(LS_WARNING) << "No selected crypto suite";
        return false;
    }
    
    RTC_LOG(LS_INFO) << "Extract DTLS-SRTP key from transport " << transport_name_;

    int key_len;
    int salt_len;
    if (!rtc::GetSrtpKeyAndSaltLengths(*selected_crypto_suite, &key_len, &salt_len)) {
        RTC_LOG(LS_WARNING) << "Unknown DTLS-SRTP crypto suite " << *selected_crypto_suite;
        return false;
    }
    
    rtc::ZeroOnFreeBuffer<unsigned char> dtls_buffer(key_len * 2 + salt_len * 2);
    if (!dtls_transport->ExportKeyingMaterial(kDtlsSrtpExporterLabel,
                NULL, 0, false, &dtls_buffer[0], dtls_buffer.size()))
    {
        RTC_LOG(LS_WARNING) << "Extracting DTLS-SRTP param failed";
        return false;
    }
    
    rtc::ZeroOnFreeBuffer<unsigned char> client_write_key(key_len + salt_len);
    rtc::ZeroOnFreeBuffer<unsigned char> server_write_key(key_len + salt_len);
    size_t offset = 0;
    memcpy(&client_write_key[0], &dtls_buffer[offset], key_len);
    offset += key_len;
    memcpy(&server_write_key[0], &dtls_buffer[offset], key_len);
    offset += key_len;
    memcpy(&client_write_key[key_len], &dtls_buffer[offset], salt_len);
    offset += salt_len;
    memcpy(&server_write_key[key_len], &dtls_buffer[offset], salt_len);
    
    rtc::SSLRole role = dtls_transport->GetDtlsRole();
    
    if (role == rtc::SSL_SERVER) {
        *send_key = std::move(server_write_key);
        *recv_key = std::move(client_write_key);
    }
    else {
        *send_key = std::move(client_write_key);
        *recv_key = std::move(server_write_key);
    }

    return true;
}

} // namespace ice


